from pyspark.mllib.evaluation import RegressionMetrics
from pyspark.mllib.recommendation import ALS, Rating, MatrixFactorizationModel
from pyspark import SparkContext
import os

spark = SparkContext.getOrCreate()
rdd = spark.textFile("E:\python_pro\P1905\lizhi\practice\ml_lesson\ml-100k\d.data")

def map_line(line):
    data = line.split("\t")

    # return Rating(int(data[0]), int(data[1]), float(data[2]))
    return (int(data[0]), int(data[1]), float(data[2]))

rdd = rdd.map(map_line)
model = ALS.train(rdd, 50, 10)

ratingForUser = model.predictAll(rdd.map(lambda r: (r[0], r[1]))).map(lambda r: (r[0], r[1], r[2])).join(rdd.map(lambda r: (r[0], r[1], r[2])))

metrics = RegressionMetrics(ratingForUser.map(lambda r: r[1]))

model.save(spark, "./rec")

# 载入模型
model = MatrixFactorizationModel.load(spark, "./rec")

print(model.predict(889, 381))
